# coding=utf-8
# 确保代码在不同的 Python 版本之间具有兼容性，特别是 Python 2 和 Python 3 的兼容性。
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import copy
import logging
import math

from os.path import join as pjoin
# PyTorch 和 NumPy 用于模型的实现和数学运算
import torch
import torch.nn as nn
import numpy as np
from torch.nn import CrossEntropyLoss, Dropout, Softmax, Linear, Conv2d, LayerNorm
from torch.nn.modules.utils import _pair
from scipy import ndimage
from .swin_transformer_unet_skip_expand_decoder_sys import SwinTransformerSys
from .swin_transformer_unetPP import SwinTansformerUnetPP
# 用于日志记录，帮助调试和监控模型训练过程
logger = logging.getLogger(__name__)
'''这些都是可以改的.....'''
# 定义 Swin Unet 网络
class SwinUnet(nn.Module):
    def __init__(self, config, img_size=224, num_classes=9, zero_head=False, vis=False):
        super(SwinUnet, self).__init__()
        self.num_classes = num_classes
        self.zero_head = zero_head
        self.config = config
        self.swin_unet = SwinTransformerSys(
        # self.swin_unet = SwinTansformerUnetPP(
            # 图像尺寸(224*224)        
            img_size=config.DATA.IMG_SIZE,
            # 补丁尺寸(4*4)
            patch_size=config.MODEL.SWIN.PATCH_SIZE,
            # 输入通道数(1)
            in_chans=config.MODEL.SWIN.IN_CHANS,
            # 输入类别数(9)
            num_classes=self.num_classes,
            # 嵌入维度(96)
            # 该参数定义了输入图像或每个补丁经过嵌入处理后的维度。它影响到模型的表达能力，嵌入维度较大有助于提升模型的表现，但也会增加计算和存储的需求
            embed_dim=config.MODEL.SWIN.EMBED_DIM,
            # 网络深度[2 , 2 , 2 , 2]
            # 指定了每个 Transformer 层的深度（即 Transformer 的堆叠层数）,较深的网络可以学习更复杂的特征，但也需要更多的计算和内存
            depths=config.MODEL.SWIN.DEPTHS,
            # 头数([3 , 6 , 12 , 24])
            # 指定了每个 Transformer 层中注意力机制的头数。每个头独立地学习输入特征，并最终将其拼接起来。增加头数可以增强模型的表达能力，但也会增加计算复杂度。
            num_heads=config.MODEL.SWIN.NUM_HEADS,
            # 窗口大小(7)
            # 局部窗口的大小决定了每个 Transformer 层在计算自注意力时关注的区域大小,较大的窗口大小允许模型关注更广的区域，但可能会丢失细节。
            window_size=config.MODEL.SWIN.WINDOW_SIZE,
            # MLP 比例(4)
            # 该参数控制每个 Transformer 层中多层感知机（MLP）部分的维度扩展比例。它影响到模型的非线性表达能力。
            mlp_ratio=config.MODEL.SWIN.MLP_RATIO,
            # QKV 偏置(True)
            # 在计算自注意力时，Q、K、V（查询、键、值）矩阵会经过线性变换并加上偏置。该参数决定是否在计算中添加 Q、K、V 的偏置项
            qkv_bias=config.MODEL.SWIN.QKV_BIAS,
            # QK 缩放(qk_scale)
            # 该参数用来缩放计算自注意力时的查询（Q）和键（K）之间的点积。在一些变体中，QK 缩放可以帮助提升模型的稳定性。
            qk_scale=config.MODEL.SWIN.QK_SCALE,
            # Dropout 比例(0)
            # Dropout 是一种常用的正则化技术，用于防止模型过拟合。该参数定义了每个层中 Dropout 的概率，通常在训练过程中使用。
            drop_rate=config.MODEL.DROP_RATE,
            # Drop Path 比例(0.1)
            # 这是一种扩展的正则化方法，在深度神经网络中随机丢弃整个网络路径。它可以帮助提高模型的泛化能力
            drop_path_rate=config.MODEL.DROP_PATH_RATE,
            # 自适应位置编码(False)
            # 该参数决定是否使用自适应位置编码（Adaptive Position Embedding），该技术可以通过自学习的方式调整位置编码，从而提高模型对不同位置关系的表示能力。
            ape=config.MODEL.SWIN.APE,
            # 补丁归一化(True)
            # 该参数控制是否在每个补丁的嵌入中使用归一化。归一化有助于加速训练，并可能提高模型的性能。
            patch_norm=config.MODEL.SWIN.PATCH_NORM,
            # 使用检查点(False)
            # 该参数决定是否在训练过程中使用梯度检查点技术。梯度检查点可以在训练时减少内存使用，但会增加计算时间。
            use_checkpoint=config.TRAIN.USE_CHECKPOINT)

    def forward(self, x):
        # 如果是灰度图那就把它转成彩色图
        if x.size()[1] == 1:
            x = x.repeat(1,3,1,1)
        logits = self.swin_unet(x)
        return logits
    # 加载预训练权重
    def load_from(self, config):
        # 获取预训练模型路径
        pretrained_path = config.MODEL.PRETRAIN_CKPT
        if pretrained_path is not None:
            # 1.从配置文件config中获取预训练权重的路径
            print("预训练权重路径:{}".format(pretrained_path))
            device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
            # 2.加载预训练权重
            pretrained_dict = torch.load(pretrained_path, map_location=device)
            # 3.处理权重
            if "model"  not in pretrained_dict:
                print("---开始加载和处理预训练权重到当前模型---")
                # 如果权重文件不包含 model 键，手动调整键名以适配当前模型结构：
                pretrained_dict = {k[17:]:v for k,v in pretrained_dict.items()}
                for k in list(pretrained_dict.keys()):
                    if "output" in k:
                        print("删除键:{}".format(k))
                        del pretrained_dict[k]
                # 把处理好的权重加载到当前模型
                msg = self.swin_unet.load_state_dict(pretrained_dict,strict=False)
                # print(msg)
                return
            pretrained_dict = pretrained_dict['model']

            print("--- 开始加载预训练权重到 swin encoder 中---")
            # 获取当前模型的权重字典
            model_dict = self.swin_unet.state_dict()  # 获取当前模型的权重字典，用于比较与预训练权重的兼容性
            full_dict = copy.deepcopy(pretrained_dict)  # 创建预训练权重的深拷贝，避免直接修改原始数据
            # 处理预训练权重的键名，以适配当前模型结构
            # 修改层数
            for k, v in pretrained_dict.items():
                # 如果键名中包含 "layers."，则这是一个 Swin Transformer 的层。
                if "layers." in k:
                    # 从预训练模型中提取层编号，并重新计算层数（假设层编号从 3 开始递减）。
                    current_layer_num = 3-int(k[7:8])
                    # 修改键名，将 "layers." 替换为 "layers_up."，以适配当前模型结构。
                    current_k = "layers_up." + str(current_layer_num) + k[8:]
                    # 更新调整后的权重字典
                    full_dict.update({current_k:v})
            # 删除不匹配的权重
            for k in list(full_dict.keys()):
                if k in model_dict:
                    if full_dict[k].shape != model_dict[k].shape:
                        print("delete:{};shape pretrain:{};shape model:{}".format(k,v.shape,model_dict[k].shape))
                        del full_dict[k]
            # 加载处理后的权重
            msg = self.swin_unet.load_state_dict(full_dict, strict=False)
            # print(msg)
        else:
            print("未发现预训练模型")

if __name__ == "__main__":
    
    model = SwinUnet().to('cuda')
    a = torch.randn(24,3,224,224).cuda()
    out = model(a)
    print(out.size())